/*	Copyright 2007 - Xavier Baro (xbaro@cvc.uab.cat)

	This file is part of eapmlib.

    Eapmlib is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or any 
	later version.

    Eapmlib is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
#include "EvolutiveLib.h"
#include <iostream>
#include "WLEAPMDipole.h"
#include "IntSamplesSet.h"
#include "NaiveBayesModel.h"
#include "Cascade.h"
#include "XMLOpenCV.h"
#include "Detector.h"


using namespace std;
using namespace Evolutive;

void detect_and_draw( IplImage* img, CvHaarClassifierCascade* cascade)
{
	int i;
    double scale = 1.3;
    IplImage* gray =NULL;
	static CvMemStorage* storage = NULL;
    static CvScalar colors[] = 
    {
        {{0,0,255}},
        {{0,128,255}},
        {{0,255,255}},
        {{0,255,0}},
        {{255,128,0}},
        {{255,255,0}},
        {{255,0,0}},
        {{255,0,255}}
    };

	storage = cvCreateMemStorage(0);

	gray = cvCreateImage( cvSize(img->width,img->height), 8, 1 );
  
    cvCvtColor( img, gray, CV_BGR2GRAY );
    cvClearMemStorage( storage );

    if( cascade )
    {
        double t = (double)cvGetTickCount();
        CvSeq* faces = cvHaarDetectObjects( gray, cascade, storage,
                                            1.1, 2, 0/*CV_HAAR_DO_CANNY_PRUNING*/,
                                            cvSize(10, 10) );
        t = (double)cvGetTickCount() - t;
        printf( "detection time = %gms\n", t/((double)cvGetTickFrequency()*1000.) );
        for( i = 0; i < (faces ? faces->total : 0); i++ )
        {
            CvRect* r = (CvRect*)cvGetSeqElem( faces, i );
            CvPoint center;
            int radius;
            center.x = cvRound((r->x + r->width*0.5));
            center.y = cvRound((r->y + r->height*0.5));
            radius = cvRound((r->width + r->height)*0.25);
            cvCircle( img, center, radius, colors[i%8], 3, 8, 0 );
        }
    }

	cvNamedWindow( "result", 1 );
    cvShowImage( "result", img );
	cvWaitKey(0);
	cvDestroyWindow("result");
    cvReleaseImage( &gray );
}


void DetectObjects(const char* CascadeFileName,const char *ImageFileName)
{
	CvHaarClassifierCascade* Cascade = NULL;
	IplImage *Image=NULL;
	
	// Load the cascade
	Cascade = (CvHaarClassifierCascade*)cvLoad(CascadeFileName, 0, 0, 0 );

	// Load the image
	Image=cvLoadImage(ImageFileName,1);

	// Show error messages
	if(!Cascade)
		cout << " Error loading the cascade" << endl;
	if(!Image)
		cout << " Error loading the image" << endl;
	if(!Cascade || !Image)
		return;

	detect_and_draw(Image,Cascade);

	// Release used memory
	cvReleaseImage( &Image );
}

void AdaboostLearning(int argc,char* argv[])
{
	CIntSamplesSet oSamplesSet;
	CAdaBoost oBoosting;
	CWLEAPMDipole oWL;
	CNaiveBayesModel oNBM;
	CAdditiveClassEnsemble *Result=NULL,*LoadedEns=new CAdditiveClassEnsemble();
	CSamplesEvaluator oEvaluator;
	CCascade oCascade;
	CXMLOpenCV oExportXML;
	CDetector oDetector;
	CRectVector oRegions;

	// Parameters
	string SamplesPath= "../data/FaceDetect_01_Train.dat";
	string OutFileName= "LastClass.log";
	int NumABIterations=10;
	int NumEvolIters=300;

	if(argc>1)
		NumABIterations=atoi(argv[1]);
	if(argc>2)
		NumEvolIters=atoi(argv[2]);
	
	// Show initial information
	cout << "TEST ADABOOST" << endl << endl;

	// Load the samples set
	cout << "Loading samples..." << endl;
	oSamplesSet.Load(SamplesPath);
	cout << "   File: " << SamplesPath << endl;
	cout << "   Num Samples: " << oSamplesSet.GetNumSamples() << endl;
	cout << "   Size: " << oSamplesSet.GetSamplesSize().width << " x " << oSamplesSet.GetSamplesSize().height << endl << endl;

	//oSamplesSet.Show();

	// Configure the weak learner
	oWL.SetProbModel(&oNBM);
	oWL.SetSamplesSet(&oSamplesSet);
	oWL.SetNumIterations(NumEvolIters);

	// Configure the adaboost object
	oBoosting.SetWeakLearner(&oWL);
	oBoosting.SetMaxIters(NumABIterations);
#if(0)
	// Start learning
	cout << "Learning Classifier..." << endl;
	oBoosting.LearnEnsemble();
	Result=oBoosting.GetEnsemble();

	// Show end message
	cout << endl << "AdaBoost learning finshed" << endl;		
	cout << "   Num classifiers: " << Result->GetNumClassifiers() << endl;
	cout << "   Output file: " << OutFileName << endl;

	oCascade.AddClassifier(Result);
	oExportXML.ExportClassifier(&oCascade,"./cascade.xml","test");

	Result->Save("./ensfile.dat");

#endif
	LoadedEns->Load("./ensfile.dat");
	oCascade.AddClassifier(LoadedEns);
	//oExportXML.ExportClassifier(&oCascade,"./cascade.xml","test");

	cout << endl << "Data after save and load" << endl;		
	cout << "   Num classifiers: " << LoadedEns->GetNumClassifiers() << endl;

	//oCascade.AddClassifier(LoadedEns);
	oExportXML.ExportClassifier(&oCascade,"./cascade2.xml","test");		

	oEvaluator.SetSamplesSet(&oSamplesSet);
	oEvaluator.SetErrorFunction(Evolutive::ERROR_MISSCLASS);
	cout << "Classifier error: " << oEvaluator.GetScore(LoadedEns) << endl;
	cout << "Cascade error: " << oEvaluator.GetScore(&oCascade) << endl;

	//DetectObjects("./cascade.xml","../data/facetest.bmp");

	//IplImage *Image=cvLoadImage("../data/facetest.bmp",0);
	IplImage *Image=cvLoadImage("../data/rand_img.jpg",0);
	oDetector.FindRegions(&oRegions,Image,&oCascade);
	oDetector.ShowRegions(&oRegions,Image);
	cvReleaseImage(&Image);
}

int main(int argc,char* argv[]) 
{	
	CWLEAPMDipole oWL;
	CNaiveBayesModel oNBM;
	CCascade oCascade;
	CIntSamplesSet oSamplesSet;
	int NumEvolIters=300;
	string SamplesPath= "..\\data\\FaceDetect_01_Train.dat";
	string BGFile="../data/BGFile2.txt";
	CDetector oObjDetector;
	CRectVector oRegions;
	string Imtest1="../data/rand_img.jpg";
	string Imtest2="../data/facetest.bmp";
	IplImage *Image=NULL;
	string CascadeFile="cascade_Thesis_Evol.dat";
				
	try
	{
		// Load the samples
		oSamplesSet.Load(SamplesPath);
		oSamplesSet.SetDetector(&oObjDetector);
		oSamplesSet.LoadBGFiles(BGFile);

		// Configure the weak learner
		oWL.SetProbModel(&oNBM);
		oWL.SetNumIterations(NumEvolIters);
		oWL.SetSamplesSet(&oSamplesSet);

		// Learn a new cascade
		oCascade.Train(0.995,0.45,7,&oWL,dynamic_cast<CSamplesSet*>(&oSamplesSet),CascadeFile);
				
		oCascade.Load(CascadeFile);		
		oObjDetector.SetClassifier(&oCascade);

		// Test the cascade
		Image=cvLoadImage(Imtest1.data(),0);
		oObjDetector.FindRegions(&oRegions,Image,&oCascade);
		oObjDetector.ShowRegions(&oRegions,Image);
		cvReleaseImage(&Image);
		Image=cvLoadImage(Imtest2.data(),0);
		oObjDetector.FindRegions(&oRegions,Image,&oCascade);
		oObjDetector.ShowRegions(&oRegions,Image);
		cvReleaseImage(&Image);
	}
	catch(CEvolutiveLibException &e)
	{
		cout << e.ToString() << endl;	
	}
	
	return 0;
}
